
rule all:
    input:
        expand("{sample}.nonshift.vcf.gz", sample=config["sample"]),
        expand("{sample}.shifted.vcf.gz", sample=config["sample"]),
        expand("{sample}.mt.snv.vcf.gz", sample=config["sample"]),

rule haplochecker:
    input:
        nonshift_bam="{sample}.nonshift.bam",
        nonshift_ref_fasta=config["nonshift_ref_fasta"],
    output:
        contamination="{sample}.contamination.txt",
        haplochecker_out=temp(directory("{sample}.haplochecker_out")),
    params:
        qual="20",
        mapq="30",
        vaf="0.01",
    shell:
        """
        mkdir -p {output.haplochecker_out} && \
        java -jar $CONDA_PREFIX/share/mitolib/mitolib-0.1.2.jar haplochecker\
            --in {input.nonshift_bam} \
            --ref {input.nonshift_ref_fasta} \
            --out {output.haplochecker_out}\
            --QUAL {params.qual} \
            --MAPQ {params.mapq} \
            --VAF {params.vaf} && \
        cp {output.haplochecker_out}/*.contamination.txt {output.contamination}
        """

rule mutect2:
    input:
        nonshift_bam="{sample}.nonshift.bam",
        nonshift_ref_fasta=config["nonshift_ref_fasta"],
        shifted_bam="{sample}.shifted.bam",
        shifted_ref_fasta=config["shifted_ref_fasta"],
    output:
        nonshift_vcf="{sample}.nonshift.vcf.gz",
        nonshift_vcf_tbi="{sample}.nonshift.vcf.gz.tbi",
        nonshift_stats="{sample}.nonshift.vcf.gz.stats",
        shifted_vcf="{sample}.shifted.vcf.gz",
        shifted_vcf_tbi="{sample}.shifted.vcf.gz.tbi",
        shifted_stats="{sample}.shifted.vcf.gz.stats",
    params:
        mutect="-A StrandBiasBySample --mitochondria-mode --max-reads-per-alignment-start 75 --max-mnp-distance 0",
        java_options=config.get("java_options", "")
    shell:
        """
        gatk {params.java_options} Mutect2 -I {input.nonshift_bam} -R {input.nonshift_ref_fasta} -O {output.nonshift_vcf} {params.mutect} && \
        gatk {params.java_options} Mutect2 -I {input.shifted_bam} -R {input.shifted_ref_fasta} -O {output.shifted_vcf} {params.mutect}
        """

rule liftover_vcf:
    input:
        nonshift_vcf="{sample}.nonshift.vcf.gz",
        shifted_vcf="{sample}.shifted.vcf.gz",
        liftover_chain=config["liftover_chain"],
        shifted_ref_fasta=config["shifted_ref_fasta"],
    output:
        shift_back_vcf="{sample}.shift_back.vcf.gz",
        shift_back_vcf_tbi="{sample}.shift_back.vcf.gz.tbi",
        rejected_vcf="{sample}.rejected.vcf.gz",
    params:
        java_options=config.get("java_options", "")
    shell:
        """
        gatk {params.java_options} LiftoverVcf \
            -I {input.shifted_vcf} \
            -R {input.shifted_ref_fasta} \
            -C {input.liftover_chain} \
            -O {output.shift_back_vcf} \
            --REJECT {output.rejected_vcf} \
            --CREATE_INDEX true
        """

rule merge_vcfs:
    input:
        nonshift_vcf="{sample}.nonshift.vcf.gz",
        shift_back_vcf="{sample}.shift_back.vcf.gz",
    output:
        raw_vcf="{sample}.raw.vcf.gz",
        raw_vcf_tbi="{sample}.raw.vcf.gz.tbi",
    params:
        java_options=config.get("java_options", "")
    shell:
        """
        gatk MergeVcfs \
            -I {input.nonshift_vcf} \
            -I {output.shift_back_vcf} \
            -O {output.raw_vcf} \
            --CREATE_INDEX true
        """

rule merge_mutect_stats:
    input:
        nonshift_stats="{sample}.nonshift.vcf.gz.stats",
        shifted_stats="{sample}.shifted.vcf.gz.stats",
    output:
        mutect_stats="{sample}.mutect.stats",
    params:
        java_options=config.get("java_options", "")
    shell:
        """
        gatk {params.java_options} MergeMutectStats \
            --stats {input.nonshift_stats} \
            --stats {input.shifted_stats} \
            -O {output.mutect_stats}
        """

def get_minor_level(contamination_file) -> str:
    """get minor level"""
    import csv
    with open(contamination_file, "r", encoding="utf8") as reader:
        csv_reader = csv.DictReader(reader, delimiter="\t")
        for row in csv_reader:
            return row["MinorLevel"]
    raise ValueError("MinorLevel value not found")

rule filter_mutect_calls:
    input:
        raw_vcf="{sample}.raw.vcf.gz",
        nonshift_ref_fasta=config["nonshift_ref_fasta"],
        mutect_stats="{sample}.mutect.stats",
        contamination="{sample}.contamination.txt",
    output:
        filtered_vcf="{sample}.filtered.vcf.gz",
        filtered_vcf_tbi="{sample}.filtered.vcf.gz.tbi",
    params:
        minor_level=lambda wildcards: get_minor_level(wildcards.sample + ".contamination.txt"),
        java_options=config.get("java_options", "")
    shell:
        """
        gatk {params.java_options} FilterMutectCalls \
            -V {input.raw_vcf} \
            -R {input.nonshift_ref_fasta} \
            -O {output.filtered_vcf} \
            --stats {input.mutect_stats} \
            --contamination-estimate {params.minor_level} \
            --max-alt-allele-count 4 \
            --mitochondria-mode
        """

rule variant_filtration:
    input:
        filtered_vcf="{sample}.filtered.vcf.gz",
        blacklist_sites=config["blacklist_sites"],
    output:
        markfil_vcf="{sample}.markfil.vcf.gz",
        markfil_vcf_tbi="{sample}.markfil.vcf.gz.tbi",
    params:
        java_options=config.get("java_options", "")
    shell:
        """
        gatk {params.java_options} VariantFiltration \
            -V {input.filtered_vcf} \
            --mask {input.blacklist_sites} \
            -O {output.markfil_vcf} \
            --mask-name 'blacklisted_site'
        """

rule normalize_vcf:
    input:
        markfil_vcf="{sample}.markfil.vcf.gz",
        nonshift_ref_fasta=config["nonshift_ref_fasta"],
    output:
        with_vaf_vcf=temp("{sample}.with_vaf.vcf.gz"),
        with_vaf_vcf_tbi=temp("{sample}.with_vaf.vcf.gz.tbi"),
        snv_vcf="{sample}.mt.snv.vcf.gz",
        snv_vcf_tbi="{sample}.mt.snv.vcf.gz.tbi",
    shell:
        """
        bcftools +fill-tags {input.markfil_vcf} -Oz -o {output.with_vaf_vcf} -- -t VAF && \
        bcftools index --tbi {output.with_vaf_vcf} && \
        bcftools norm -f {input.nonshift_ref_fasta} -m -any {output.with_vaf_vcf} | \
        awk '$1~/#/ || $5!~/[^ATGC]/' | bgzip -c > {output.snv_vcf} && \
        bcftools index --tbi {output.snv_vcf}
        """
